package io.prometheus.jmx;

import io.prometheus.client.Collector;
import io.prometheus.client.Counter;
import org.yaml.snakeyaml.Yaml;

import javax.management.MalformedObjectNameException;
import javax.management.ObjectName;
import java.io.File;
import java.io.FileReader;
import java.io.IOException;
import java.io.PrintWriter;
import java.io.StringWriter;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.TreeMap;
import java.util.logging.Logger;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import static java.lang.String.format;

public class JmxCollector extends Collector implements Collector.Describable {
  static final Counter configReloadSuccess = Counter.build()
          .name("jmx_config_reload_success_total")
          .help("Number of times configuration have successfully been reloaded.").register();

  static final Counter configReloadFailure = Counter.build()
          .name("jmx_config_reload_failure_total")
          .help("Number of times configuration have failed to be reloaded.").register();

  private static final Logger LOGGER = Logger.getLogger(JmxCollector.class.getName());

  private static class Rule {
    Pattern pattern;
    String name;
    String value;
    Double valueFactor = 1.0;
    String help;
    boolean attrNameSnakeCase;
    Type type = Type.UNTYPED;
    ArrayList<String> labelNames;
    ArrayList<String> labelValues;
  }

  private static class Config {
    Integer startDelaySeconds = 0;
    String jmxUrl = "";
    String username = "";
    String password = "";
    boolean ssl = false;
    boolean lowercaseOutputName;
    boolean lowercaseOutputLabelNames;
    List<ObjectName> whitelistObjectNames = new ArrayList<ObjectName>();
    List<ObjectName> blacklistObjectNames = new ArrayList<ObjectName>();
    List<Rule> rules = new ArrayList<Rule>();
    long lastUpdate = 0L;
  }

  private Config config;
  //private File configFile;
  private long createTimeNanoSecs = System.nanoTime();

  private final JmxMBeanPropertyCache jmxMBeanPropertyCache = new JmxMBeanPropertyCache();

  public JmxCollector(Map<String, Object> map) throws IOException, MalformedObjectNameException {
    //configFile = in;
    config = loadConfig(map);
    //config.lastUpdate = configFile.lastModified();
  }

  /*public JmxCollector(String yamlConfig) throws MalformedObjectNameException {
    config = loadConfig((Map<String, Object>)new Yaml().load(yamlConfig));
  }*/

 /* private void reloadConfig() {
    try {
      FileReader fr = new FileReader(configFile);

      try {
        Map<String, Object> newYamlConfig = (Map<String, Object>)new Yaml().load(fr);
        config = loadConfig(newYamlConfig);
        config.lastUpdate = configFile.lastModified();
        configReloadSuccess.inc();
      } catch (Exception e) {
        LOGGER.severe("Configuration reload failed: " + e.toString());
        configReloadFailure.inc();
      } finally {
        fr.close();
      }

    } catch (IOException e) {
      LOGGER.severe("Configuration reload failed: " + e.toString());
      configReloadFailure.inc();
    }
  }*/

  private Config loadConfig(Map<String, Object> yamlConfig) throws MalformedObjectNameException {
    Config cfg = new Config();

    if (yamlConfig == null) {  // Yaml config empty, set config to empty map.
      yamlConfig = new HashMap<String, Object>();
    }

    if (yamlConfig.containsKey("startDelaySeconds")) {
      try {
        cfg.startDelaySeconds = (Integer) yamlConfig.get("startDelaySeconds");
      } catch (NumberFormatException e) {
        throw new IllegalArgumentException("Invalid number provided for startDelaySeconds", e);
      }
    }
    if (yamlConfig.containsKey("hostPort")) {
      if (yamlConfig.containsKey("jmxUrl")) {
        throw new IllegalArgumentException("At most one of hostPort and jmxUrl must be provided");
      }
      cfg.jmxUrl ="service:jmx:rmi:///jndi/rmi://" + (String)yamlConfig.get("hostPort") + "/jmxrmi";
    } else if (yamlConfig.containsKey("jmxUrl")) {
      cfg.jmxUrl = (String)yamlConfig.get("jmxUrl");
    }

    if (yamlConfig.containsKey("username")) {
      cfg.username = (String)yamlConfig.get("username");
    }

    if (yamlConfig.containsKey("password")) {
      cfg.password = (String)yamlConfig.get("password");
    }

    if (yamlConfig.containsKey("ssl")) {
      cfg.ssl = (Boolean)yamlConfig.get("ssl");
    }

    if (yamlConfig.containsKey("lowercaseOutputName")) {
      cfg.lowercaseOutputName = (Boolean)yamlConfig.get("lowercaseOutputName");
    }

    if (yamlConfig.containsKey("lowercaseOutputLabelNames")) {
      cfg.lowercaseOutputLabelNames = (Boolean)yamlConfig.get("lowercaseOutputLabelNames");
    }

    if (yamlConfig.containsKey("whitelistObjectNames")) {
      List<Object> names = (List<Object>) yamlConfig.get("whitelistObjectNames");
      for(Object name : names) {
        cfg.whitelistObjectNames.add(new ObjectName((String)name));
      }
    } else {
      cfg.whitelistObjectNames.add(null);
    }

    if (yamlConfig.containsKey("blacklistObjectNames")) {
      List<Object> names = (List<Object>) yamlConfig.get("blacklistObjectNames");
      for (Object name : names) {
        cfg.blacklistObjectNames.add(new ObjectName((String)name));
      }
    }

    if (yamlConfig.containsKey("rules")) {
      List<Map<String,Object>> configRules = (List<Map<String,Object>>) yamlConfig.get("rules");
      for (Map<String, Object> ruleObject : configRules) {
        Map<String, Object> yamlRule = ruleObject;
        Rule rule = new Rule();
        cfg.rules.add(rule);
        if (yamlRule.containsKey("pattern")) {
          rule.pattern = Pattern.compile("^.*(?:" + (String)yamlRule.get("pattern") + ").*$");
        }
        if (yamlRule.containsKey("name")) {
          rule.name = (String)yamlRule.get("name");
        }
        if (yamlRule.containsKey("value")) {
          rule.value = String.valueOf(yamlRule.get("value"));
        }
        if (yamlRule.containsKey("valueFactor")) {
          String valueFactor = String.valueOf(yamlRule.get("valueFactor"));
          try {
            rule.valueFactor = Double.valueOf(valueFactor);
          } catch (NumberFormatException e) {
            // use default value
          }
        }
        if (yamlRule.containsKey("attrNameSnakeCase")) {
          rule.attrNameSnakeCase = (Boolean)yamlRule.get("attrNameSnakeCase");
        }
        if (yamlRule.containsKey("type")) {
          rule.type = Type.valueOf((String)yamlRule.get("type"));
        }
        if (yamlRule.containsKey("help")) {
          rule.help = (String)yamlRule.get("help");
        }
        if (yamlRule.containsKey("labels")) {
          TreeMap labels = new TreeMap((Map<String, Object>)yamlRule.get("labels"));
          rule.labelNames = new ArrayList<String>();
          rule.labelValues = new ArrayList<String>();
          for (Map.Entry<String, Object> entry : (Set<Map.Entry<String, Object>>)labels.entrySet()) {
            rule.labelNames.add(entry.getKey());
            rule.labelValues.add((String)entry.getValue());
          }
        }

        // Validation.
        if ((rule.labelNames != null || rule.help != null) && rule.name == null) {
          throw new IllegalArgumentException("Must provide name, if help or labels are given: " + yamlRule);
        }
        if (rule.name != null && rule.pattern == null) {
          throw new IllegalArgumentException("Must provide pattern, if name is given: " + yamlRule);
        }
      }
    } else {
      // Default to a single default rule.
      cfg.rules.add(new Rule());
    }

    return cfg;

  }

  static String toSnakeAndLowerCase(String attrName) {
    if (attrName == null || attrName.isEmpty()) {
      return attrName;
    }
    char firstChar = attrName.subSequence(0, 1).charAt(0);
    boolean prevCharIsUpperCaseOrUnderscore = Character.isUpperCase(firstChar) || firstChar == '_';
    StringBuilder resultBuilder = new StringBuilder(attrName.length()).append(Character.toLowerCase(firstChar));
    for (char attrChar : attrName.substring(1).toCharArray()) {
      boolean charIsUpperCase = Character.isUpperCase(attrChar);
      if (!prevCharIsUpperCaseOrUnderscore && charIsUpperCase) {
        resultBuilder.append("_");
      }
      resultBuilder.append(Character.toLowerCase(attrChar));
      prevCharIsUpperCaseOrUnderscore = charIsUpperCase || attrChar == '_';
    }
    return resultBuilder.toString();
  }

  /**
   * Change invalid chars to underscore, and merge underscores.
   * @param name Input string
   * @return
   */
  static String safeName(String name) {
    if (name == null) {
      return null;
    }
    boolean prevCharIsUnderscore = false;
    StringBuilder safeNameBuilder = new StringBuilder(name.length());
    if (!name.isEmpty() && Character.isDigit(name.charAt(0))) {
      // prevent a numeric prefix.
      safeNameBuilder.append("_");
    }
    for (char nameChar : name.toCharArray()) {
      boolean isUnsafeChar = !JmxCollector.isLegalCharacter(nameChar);
      if ((isUnsafeChar || nameChar == '_')) {
        if (prevCharIsUnderscore) {
          continue;
        } else {
          safeNameBuilder.append("_");
          prevCharIsUnderscore = true;
        }
      } else {
        safeNameBuilder.append(nameChar);
        prevCharIsUnderscore = false;
      }
    }

    return safeNameBuilder.toString();
  }

  private static boolean isLegalCharacter(char input) {
    return ((input == ':') ||
            (input == '_') ||
            (input >= 'a' && input <= 'z') ||
            (input >= 'A' && input <= 'Z') ||
            (input >= '0' && input <= '9'));
  }

  class Receiver implements JmxScraper.MBeanReceiver {
    Map<String, MetricFamilySamples> metricFamilySamplesMap =
            new HashMap<String, MetricFamilySamples>();

    private static final char SEP = '_';



    // [] and () are special in regexes, so swtich to <>.
    private String angleBrackets(String s) {
      return "<" + s.substring(1, s.length() - 1) + ">";
    }

    void addSample(MetricFamilySamples.Sample sample, Type type, String help) {
      MetricFamilySamples mfs = metricFamilySamplesMap.get(sample.name);
      if (mfs == null) {
        // JmxScraper.MBeanReceiver is only called from one thread,
        // so there's no race here.
        mfs = new MetricFamilySamples(sample.name, type, help, new ArrayList<MetricFamilySamples.Sample>());
        metricFamilySamplesMap.put(sample.name, mfs);
      }
      mfs.samples.add(sample);
    }

    private void defaultExport(
            String domain,
            LinkedHashMap<String, String> beanProperties,
            LinkedList<String> attrKeys,
            String attrName,
            String help,
            Object value,
            Type type) {
      StringBuilder name = new StringBuilder();
      name.append(domain);
      if (beanProperties.size() > 0) {
        name.append(SEP);
        name.append(beanProperties.values().iterator().next());
      }
      for (String k : attrKeys) {
        name.append(SEP);
        name.append(k);
      }
      name.append(SEP);
      name.append(attrName);
      String fullname = safeName(name.toString());

      if (config.lowercaseOutputName) {
        fullname = fullname.toLowerCase();
      }

      List<String> labelNames = new ArrayList<String>();
      List<String> labelValues = new ArrayList<String>();
      if (beanProperties.size() > 1) {
        Iterator<Map.Entry<String, String>> iter = beanProperties.entrySet().iterator();
        // Skip the first one, it's been used in the name.
        iter.next();
        while (iter.hasNext()) {
          Map.Entry<String, String> entry = iter.next();
          String labelName = safeName(entry.getKey());
          if (config.lowercaseOutputLabelNames) {
            labelName = labelName.toLowerCase();
          }
          labelNames.add(labelName);
          labelValues.add(entry.getValue());
        }
      }

      addSample(new MetricFamilySamples.Sample(fullname, labelNames, labelValues, ((Number)value).doubleValue()),
              type, help);
    }

    public void recordBean(
            String domain,
            LinkedHashMap<String, String> beanProperties,
            LinkedList<String> attrKeys,
            String attrName,
            String attrType,
            String attrDescription,
            Object beanValue) {

      String beanName = domain + angleBrackets(beanProperties.toString()) + angleBrackets(attrKeys.toString());
      // attrDescription tends not to be useful, so give the fully qualified name too.
      String help = attrDescription + " (" + beanName + attrName + ")";
      String attrNameSnakeCase = toSnakeAndLowerCase(attrName);

      for (Rule rule : config.rules) {
        Matcher matcher = null;
        String matchName = beanName + (rule.attrNameSnakeCase ? attrNameSnakeCase : attrName);
        if (rule.pattern != null) {
          matcher = rule.pattern.matcher(matchName + ": " + beanValue);
          if (!matcher.matches()) {
            continue;
          }
        }

        Number value;
        if (rule.value != null && !rule.value.isEmpty()) {
          String val = matcher.replaceAll(rule.value);

          try {
            beanValue = Double.valueOf(val);
          } catch (NumberFormatException e) {
            LOGGER.fine("Unable to parse configured value '" + val + "' to number for bean: " + beanName + attrName + ": " + beanValue);
            return;
          }
        }
        if (beanValue instanceof Number) {
          value = ((Number)beanValue).doubleValue() * rule.valueFactor;
        } else if (beanValue instanceof Boolean) {
          value = (Boolean)beanValue ? 1 : 0;
        } else {
          LOGGER.fine("Ignoring unsupported bean: " + beanName + attrName + ": " + beanValue);
          return;
        }

        // If there's no name provided, use default export format.
        if (rule.name == null) {
          defaultExport(domain, beanProperties, attrKeys, rule.attrNameSnakeCase ? attrNameSnakeCase : attrName, help, value, rule.type);
          return;
        }

        // Matcher is set below here due to validation in the constructor.
        String name = safeName(matcher.replaceAll(rule.name));
        if (name.isEmpty()) {
          return;
        }
        if (config.lowercaseOutputName) {
          name = name.toLowerCase();
        }

        // Set the help.
        if (rule.help != null) {
          help = matcher.replaceAll(rule.help);
        }

        // Set the labels.
        ArrayList<String> labelNames = new ArrayList<String>();
        ArrayList<String> labelValues = new ArrayList<String>();
        if (rule.labelNames != null) {
          for (int i = 0; i < rule.labelNames.size(); i++) {
            final String unsafeLabelName = rule.labelNames.get(i);
            final String labelValReplacement = rule.labelValues.get(i);
            try {
              String labelName = safeName(matcher.replaceAll(unsafeLabelName));
              String labelValue = matcher.replaceAll(labelValReplacement);
              if (config.lowercaseOutputLabelNames) {
                labelName = labelName.toLowerCase();
              }
              if (!labelName.isEmpty() && !labelValue.isEmpty()) {
                labelNames.add(labelName);
                labelValues.add(labelValue);
              }
            } catch (Exception e) {
              throw new RuntimeException(
                      format("Matcher '%s' unable to use: '%s' value: '%s'", matcher, unsafeLabelName, labelValReplacement), e);
            }
          }
        }

        // Add to samples.
        LOGGER.fine("add metric sample: " + name + " " + labelNames + " " + labelValues + " " + value.doubleValue());
        addSample(new MetricFamilySamples.Sample(name, labelNames, labelValues, value.doubleValue()), rule.type, help);
        return;
      }
    }

  }

  public List<MetricFamilySamples> collect() {
    /*if (configFile != null) {
      long mtime = configFile.lastModified();
      if (mtime > config.lastUpdate) {
        LOGGER.fine("Configuration file changed, reloading...");
        reloadConfig();
      }
    }*/

    Receiver receiver = new Receiver();
    JmxScraper scraper = new JmxScraper(config.jmxUrl, config.username, config.password, config.ssl,
            config.whitelistObjectNames, config.blacklistObjectNames, receiver, jmxMBeanPropertyCache);
    long start = System.nanoTime();
    double error = 0;
    if ((config.startDelaySeconds > 0) &&
            ((start - createTimeNanoSecs) / 1000000000L < config.startDelaySeconds)) {
      throw new IllegalStateException("JMXCollector waiting for startDelaySeconds");
    }
    try {
      scraper.doScrape();
    } catch (Exception e) {
      error = 1;
      StringWriter sw = new StringWriter();
      e.printStackTrace(new PrintWriter(sw));
      LOGGER.severe("JMX scrape failed: " + sw.toString());
    }
    List<MetricFamilySamples> mfsList = new ArrayList<MetricFamilySamples>();
    mfsList.addAll(receiver.metricFamilySamplesMap.values());
    List<MetricFamilySamples.Sample> samples = new ArrayList<MetricFamilySamples.Sample>();
    samples.add(new MetricFamilySamples.Sample(
            "jmx_scrape_duration_seconds", new ArrayList<String>(), new ArrayList<String>(), (System.nanoTime() - start) / 1.0E9));
    mfsList.add(new MetricFamilySamples("jmx_scrape_duration_seconds", Type.GAUGE, "Time this JMX scrape took, in seconds.", samples));

    samples = new ArrayList<MetricFamilySamples.Sample>();
    samples.add(new MetricFamilySamples.Sample(
            "jmx_scrape_error", new ArrayList<String>(), new ArrayList<String>(), error));
    mfsList.add(new MetricFamilySamples("jmx_scrape_error", Type.GAUGE, "Non-zero if this scrape failed.", samples));
    return mfsList;
  }

  public List<MetricFamilySamples> describe() {
    List<MetricFamilySamples> sampleFamilies = new ArrayList<MetricFamilySamples>();
    sampleFamilies.add(new MetricFamilySamples("jmx_scrape_duration_seconds", Type.GAUGE, "Time this JMX scrape took, in seconds.", new ArrayList<MetricFamilySamples.Sample>()));
    sampleFamilies.add(new MetricFamilySamples("jmx_scrape_error", Type.GAUGE, "Non-zero if this scrape failed.", new ArrayList<MetricFamilySamples.Sample>()));
    return sampleFamilies;
  }

  /**
   * Convenience function to run standalone.
   */
  public static void main(String[] args) throws Exception {
    String hostPort = "";
    if (args.length > 0) {
      hostPort = args[0];
    }
    /*JmxCollector jc = new JmxCollector(("{"
            + "`hostPort`: `" + hostPort + "`,"
            + "}").replace('`', '"'));
    for(MetricFamilySamples mfs : jc.collect()) {
      System.out.println(mfs);
    }*/
  }
}